#include "mcp_capability/ToolsCapability.h"
#include "mcp_capability/CapabilityMethod.h"
#include "mock/MockMessageSerializer.h"
#include <gtest/gtest.h>
#include <QJsonObject>
#include <QJsonArray>
#include <QJsonDocument>
#include <memory>

using ::testing::Return;
using ::testing::NiceMock;

class ToolsCapabilityTest : public ::testing::Test
{
protected:
    void SetUp() override
    {
        toolsCapability = new ToolsCapability();
    }
    
    void TearDown() override
    {
        delete toolsCapability;
    }
    
    ToolsCapability* toolsCapability;
};

// Test constructor and basic properties
TEST_F(ToolsCapabilityTest, ConstructorAndPrefix)
{
    EXPECT_EQ(toolsCapability->prefix(), "tools");
    
    QJsonObject capability = toolsCapability->capability();
    EXPECT_TRUE(capability.contains("listChanged"));
    EXPECT_FALSE(capability["listChanged"].toBool());
}

// Test setCapability method
TEST_F(ToolsCapabilityTest, SetCapabilityValid)
{
    QString validCapability = "listChanged";
    
    toolsCapability->setCapability(validCapability, true);
    QJsonObject capability = toolsCapability->capability();
    
    EXPECT_TRUE(capability.contains("listChanged"));
    EXPECT_TRUE(capability["listChanged"].toBool());
}

TEST_F(ToolsCapabilityTest, SetCapabilityInvalid)
{
    QString invalidCapability = "invalid json";
    
    // Should not crash and should keep previous capability
    QJsonObject originalCapability = toolsCapability->capability();
    toolsCapability->setCapability(invalidCapability);
    QJsonObject currentCapability = toolsCapability->capability();
    
    // Should remain unchanged
    EXPECT_EQ(originalCapability["name"].toString(), currentCapability["name"].toString());
}

// Test addMethod validation
TEST_F(ToolsCapabilityTest, AddMethodValid)
{
    QJsonObject methodDesc{
        {"name", "test_tool"},
        {"description", "A test tool"},
        {"inputSchema", QJsonObject{
            {"type", "object"},
            {"properties", QJsonObject{
                {"param1", QJsonObject{
                    {"type", "string"}
                }}
            }}
        }}
    };
    
    auto methodTask = createTask([](const QJsonObject&)->QJsonObject
        {return QJsonObject();}
    );

    CapabilityMethod validMethod{methodDesc, methodTask};
    bool result = toolsCapability->addMethod(validMethod);
    EXPECT_TRUE(result);
    
    QJsonArray methods = toolsCapability->list();
    EXPECT_EQ(methods.size(), 1);
    EXPECT_EQ(methods[0].toObject()["name"].toString(), "test_tool");
}

TEST_F(ToolsCapabilityTest, AddMethodMissingName)
{
    QJsonObject methodDesc{
        {"description", "A tool without name"}
    };
    
    CapabilityMethod invalidMethod{methodDesc};
    bool result = toolsCapability->addMethod(invalidMethod);
    EXPECT_FALSE(result);
    
    QJsonArray methods = toolsCapability->list();
    EXPECT_EQ(methods.size(), 0);
}

TEST_F(ToolsCapabilityTest, AddMethodMissingCallbackFunction)
{
    QJsonObject methodDesc{
        {"name", "no_description_tool"}
    };
    
    CapabilityMethod invalidMethod{methodDesc};
    bool result = toolsCapability->addMethod(invalidMethod);
    EXPECT_FALSE(result);
    
    QJsonArray methods = toolsCapability->list();
    EXPECT_EQ(methods.size(), 0);
}

TEST_F(ToolsCapabilityTest, AddMethodInvalidParameters)
{
    QJsonObject methodDesc{
        {"name", "invalid_params_tool"},
        {"description", "Tool with invalid parameters"},
        {"inputSchema", "should_be_object_not_string"}
    };
    
    CapabilityMethod invalidMethod{methodDesc};
    bool result = toolsCapability->addMethod(invalidMethod);
    EXPECT_FALSE(result);
    
    QJsonArray methods = toolsCapability->list();
    EXPECT_EQ(methods.size(), 0);
}

// Test multiple methods
TEST_F(ToolsCapabilityTest, AddMultipleMethods)
{
    QJsonObject methodDesc1{
        {"name", "tool1"},
        {"description", "First tool"}
    };
    
    auto methodTask1 = createTask([](const QJsonObject&)->QJsonObject
        {return QJsonObject{{"data", "first tool result"}};}
    );

    QJsonObject methodDesc2{
        {"name", "tool2"},
        {"description", "Second tool"}
    };
    
    auto methodTask2 = createTask([](const QJsonObject&)->QJsonObject
        {return QJsonObject{{"data", "second tool result"}};}
    );

    CapabilityMethod method1{methodDesc1, methodTask1};
    CapabilityMethod method2{methodDesc2, methodTask2};
    
    EXPECT_TRUE(toolsCapability->addMethod(method1));
    EXPECT_TRUE(toolsCapability->addMethod(method2));
    
    QJsonArray methods = toolsCapability->list();
    EXPECT_EQ(methods.size(), 2);
    
    // Check that both methods are present
    bool foundTool1 = false;
    bool foundTool2 = false;
    
    for (const auto& method : qAsConst(methods)) {
        QString name = method.toObject()["name"].toString();
        if (name == "tool1") foundTool1 = true;
        if (name == "tool2") foundTool2 = true;
    }
    
    EXPECT_TRUE(foundTool1);
    EXPECT_TRUE(foundTool2);
}

// Test clearMethod
TEST_F(ToolsCapabilityTest, ClearMethod)
{
    QJsonObject methodDesc{
        {"name", "to_clear"},
        {"description", "Will be cleared"}
    };

    auto methodTask = createTask([](const QJsonObject&)->QJsonObject
        {return QJsonObject();}
    );

    CapabilityMethod validMethod{methodDesc, methodTask};
    EXPECT_TRUE(toolsCapability->addMethod(validMethod));
    EXPECT_EQ(toolsCapability->list().size(), 1);
    
    toolsCapability->clearMethod();
    EXPECT_EQ(toolsCapability->list().size(), 0);
}


// Test list method with empty state
TEST_F(ToolsCapabilityTest, ListEmpty)
{
    QJsonArray methods = toolsCapability->list();
    EXPECT_TRUE(methods.isEmpty());
}

// Test complex method validation
TEST_F(ToolsCapabilityTest, ComplexMethodValidation)
{
    QJsonObject methodDesc{
        {"name", "complex_tool"},
        {"description", "A complex tool with full schema"},
        {"inputSchema", QJsonObject{
            {"type", "object"},
            {"properties", QJsonObject{
                {"stringParam", QJsonObject{
                    {"type", "string"},
                    {"description", "A string parameter"}
                }},
                {"numberParam", QJsonObject{
                    {"type", "number"},
                    {"minimum", 0}
                }}
            }},
            {"required", QJsonArray{"stringParam"}}
        }}
    };
    
    auto methodTask = createTask([](const QJsonObject&)->QJsonObject
        {return QJsonObject();}
    );

    CapabilityMethod complexMethod{methodDesc, methodTask};
    bool result = toolsCapability->addMethod(complexMethod);
    EXPECT_TRUE(result);
    
    QJsonArray methods = toolsCapability->list();
    EXPECT_EQ(methods.size(), 1);
    
    QJsonObject storedMethod = methods[0].toObject();
    EXPECT_EQ(storedMethod["name"].toString(), "complex_tool");
    EXPECT_TRUE(storedMethod.contains("inputSchema"));
}

// Test doRequest with tools/list method
TEST_F(ToolsCapabilityTest, DoRequestToolsList)
{
    // Add a test method first
    QJsonObject methodDesc{
        {"name", "test_tool"},
        {"description", "A test tool"}
    };
    
    auto methodTask = createTask([](const QJsonObject&)->QJsonObject
        {return QJsonObject();}
    );

    CapabilityMethod testMethod{methodDesc, methodTask};
    EXPECT_TRUE(toolsCapability->addMethod(testMethod));
    
    // Create tools/list request
    QJsonObject request{
        {"jsonrpc", "2.0"},
        {"id", 1},
        {"method", "tools/list"}
    };
    
    // Set up mock expectations
    QJsonObject expectedResult{{"tools", toolsCapability->list()}};
    QJsonObject expectedResponse;
    
    auto mockSerializer = std::make_unique<NiceMock<MockMessageSerializer>>();
    EXPECT_CALL(*mockSerializer, serializeSuccessResponse(QVariant(1), expectedResult))
        .WillOnce(::testing::Return(expectedResponse));

    toolsCapability->setSerializer(std::move(mockSerializer));
    QJsonObject response = toolsCapability->doRequest(request);
    
    // Verify that the mock was called
    EXPECT_EQ(response, expectedResponse);
}

// Test doRequest with tools/call method - successful execution
TEST_F(ToolsCapabilityTest, DoRequestToolsCallSuccess)
{
    // Add a test method with callback
    QJsonObject methodDesc{
        {"name", "add_numbers"},
        {"description", "Add two numbers"},
        {"inputSchema", QJsonObject{
            {"type", "object"},
            {"properties", QJsonObject{
                {"a", QJsonObject{{"type", "number"}}},
                {"b", QJsonObject{{"type", "number"}}}
            }},
            {"required", QJsonArray{"a", "b"}}
        }}
    };
    
    auto methodTask = createTask([](const QJsonObject& params)->QJsonObject
    {
        double a = params["a"].toDouble();
        double b = params["b"].toDouble();
        return QJsonObject{{"sum", a + b}};
    });

    CapabilityMethod testMethod{methodDesc, methodTask};
    EXPECT_TRUE(toolsCapability->addMethod(testMethod));
    
    // Create tools/call request
    QJsonObject request{
        {"jsonrpc", "2.0"},
        {"id", "test_call"},
        {"method", "tools/call"},
        {"params", QJsonObject{
            {"name", "add_numbers"},
            {"arguments", QJsonObject{
                {"a", 5},
                {"b", 3}
            }}
        }}
    };
    
    // Set up mock expectations
    QJsonObject expectedResult{{"sum", 8.0}};
    QJsonObject expectedResponse;
    
    auto mockSerializer = std::make_unique<NiceMock<MockMessageSerializer>>();
    EXPECT_CALL(*mockSerializer, serializeSuccessResponse(QVariant("test_call"), expectedResult))
        .WillOnce(::testing::Return(expectedResponse));
    
    toolsCapability->setSerializer(std::move(mockSerializer));
    QJsonObject response = toolsCapability->doRequest(request);
    
    // Verify that the mock was called
    EXPECT_EQ(response, expectedResponse);
}

// Test doRequest with tools/call method - unknown tool
TEST_F(ToolsCapabilityTest, DoRequestToolsCallUnknownTool)
{
    // Create tools/call request for non-existent tool
    QJsonObject request{
        {"jsonrpc", "2.0"},
        {"id", 123},
        {"method", "tools/call"},
        {"params", QJsonObject{
            {"name", "unknown_tool"},
            {"arguments", QJsonObject{}}
        }}
    };
    
    // Set up mock expectations - only verify that the method is called with correct parameters
    auto mockSerializer = std::make_unique<NiceMock<MockMessageSerializer>>();
    EXPECT_CALL(*mockSerializer, serializeErrorResponse(QVariant(123), -32602, ::testing::_))
        .WillOnce(::testing::Return(QJsonObject{}));
    
    toolsCapability->setSerializer(std::move(mockSerializer));
    toolsCapability->doRequest(request);
}

// Test doRequest with tools/call method - missing required parameters
TEST_F(ToolsCapabilityTest, DoRequestToolsCallMissingRequiredParams)
{
    // Add a test method with required parameters
    QJsonObject methodDesc{
        {"name", "test_required"},
        {"description", "Test required parameters"},
        {"inputSchema", QJsonObject{
            {"type", "object"},
            {"properties", QJsonObject{
                {"required_param", QJsonObject{{"type", "string"}}}
            }},
            {"required", QJsonArray{"required_param"}}
        }}
    };
    
    auto methodTask = createTask([](const QJsonObject&)->QJsonObject
        {return QJsonObject();}
    );

    CapabilityMethod testMethod{methodDesc, methodTask};
    EXPECT_TRUE(toolsCapability->addMethod(testMethod));
    
    // Create tools/call request without required parameter
    QJsonObject request{
        {"jsonrpc", "2.0"},
        {"id", "missing_params"},
        {"method", "tools/call"},
        {"params", QJsonObject{
            {"name", "test_required"},
            {"arguments", QJsonObject{}} // Missing required_param
        }}
    };
    
    // Set up mock expectations - only verify that the method is called with correct parameters
    auto mockSerializer = std::make_unique<NiceMock<MockMessageSerializer>>();
    EXPECT_CALL(*mockSerializer, serializeErrorResponse(QVariant("missing_params"), -32602, ::testing::_))
        .WillOnce(::testing::Return(QJsonObject{}));
    
    toolsCapability->setSerializer(std::move(mockSerializer));
    toolsCapability->doRequest(request);
}

// Test doRequest with unknown method
TEST_F(ToolsCapabilityTest, DoRequestUnknownMethod)
{
    // Create request for unknown method
    QJsonObject request{
        {"jsonrpc", "2.0"},
        {"id", 999},
        {"method", "unknown_method"},
        {"params", QJsonObject{}}
    };
    
    // Set up mock expectations - only verify that the method is called with correct parameters
    auto mockSerializer = std::make_unique<NiceMock<MockMessageSerializer>>();
    EXPECT_CALL(*mockSerializer, serializeErrorResponse(QVariant(999), -32601, ::testing::_))
        .WillOnce(::testing::Return(QJsonObject{}));
    
    toolsCapability->setSerializer(std::move(mockSerializer));
    toolsCapability->doRequest(request);
}

// Test doRequest with empty tools list
TEST_F(ToolsCapabilityTest, DoRequestToolsListEmpty)
{
    // Ensure no methods are added
    toolsCapability->clearMethod();
    
    // Create tools/list request
    QJsonObject request{
        {"jsonrpc", "2.0"},
        {"id", "empty_list"},
        {"method", "tools/list"}
    };
    
    QJsonObject response = toolsCapability->doRequest(request);
    
    // Verify response structure
    EXPECT_TRUE(response.contains("jsonrpc"));
    EXPECT_EQ(response["jsonrpc"].toString(), "2.0");
    EXPECT_TRUE(response.contains("id"));
    EXPECT_EQ(response["id"].toString(), "empty_list");
    EXPECT_TRUE(response.contains("result"));
    
    QJsonObject result = response["result"].toObject();
    EXPECT_TRUE(result.contains("tools"));
    
    QJsonArray tools = result["tools"].toArray();
    EXPECT_TRUE(tools.isEmpty());
}
